package DataEncryptionStandard;

public class DES {

    private final static int[] IP = new int[]{
            58, 50, 42, 34, 26, 18, 10, 2,
            60, 52, 44, 36, 28, 20, 12, 4,
            62, 54, 46, 38, 30, 22, 14, 6,
            64, 56, 48, 40, 32, 24, 16, 8,
            57, 49, 41, 33, 25, 17, 9, 1,
            59, 51, 43, 35, 27, 19, 11, 3,
            61, 53, 45, 37, 29, 21, 13, 5,
            63, 55, 47, 39, 31, 23, 15, 7
    };

    private final static int[] IPInverse = new int[]{
            40, 8, 48, 16, 56, 24, 64, 32,
            39, 7, 47, 15, 55, 23, 63, 31,
            38, 6, 46, 14, 54, 22, 62, 30,
            37, 5, 45, 13, 53, 21, 61, 29,
            36, 4, 44, 12, 52, 20, 60, 28,
            35, 3, 43, 11, 51, 19, 59, 27,
            34, 2, 42, 10, 50, 18, 58, 26,
            33, 1, 41, 9, 49, 17, 57, 25
    };

    private final static int[] EXTEND = new int[]{
            32, 1, 2, 3, 4, 5,
            4, 5, 6, 7, 8, 9,
            8, 9, 10, 11, 12, 13,
            12, 13, 14, 15, 16, 17,
            16, 17, 18, 19, 20, 21,
            20, 21, 22, 23, 24, 25,
            24, 25, 26, 27, 28, 29,
            28, 29, 30, 31, 32, 1
    };

    private final static int[][][] SBOXs = new int[][][]{
            {
                    {14, 4, 13, 1, 2, 15, 11, 8, 3, 10, 6, 12, 5, 9, 0, 7},
                    {0, 15, 7, 4, 14, 2, 13, 1, 10, 6, 12, 11, 9, 5, 3, 8},
                    {4, 1, 14, 8, 13, 6, 2, 11, 15, 12, 9, 7, 3, 10, 5, 0},
                    {15, 12, 8, 2, 4, 9, 1, 7, 5, 11, 3, 14, 10, 0, 6, 13}
            },
            {
                    {15, 1, 8, 14, 6, 11, 3, 4, 9, 7, 2, 13, 12, 0, 5, 10},
                    {3, 13, 4, 7, 15, 2, 8, 14, 12, 0, 1, 10, 6, 9, 11, 5},
                    {0, 14, 7, 11, 10, 4, 13, 1, 5, 8, 12, 6, 9, 3, 2, 15},
                    {13, 8, 10, 1, 3, 15, 4, 2, 11, 6, 7, 12, 0, 5, 14, 9}
            },
            {
                    {10, 0, 9, 14, 6, 3, 15, 5, 1, 13, 12, 7, 11, 4, 2, 8},
                    {13, 7, 0, 9, 3, 4, 6, 10, 2, 8, 5, 14, 12, 11, 15, 1},
                    {13, 6, 4, 9, 8, 15, 3, 0, 11, 1, 2, 12, 5, 10, 14, 7},
                    {1, 10, 13, 0, 6, 9, 8, 7, 4, 15, 14, 3, 11, 5, 2, 12}
            },
            {
                    {7, 13, 14, 3, 0, 6, 9, 10, 1, 2, 8, 5, 11, 12, 4, 15},
                    {13, 8, 11, 5, 6, 15, 0, 3, 4, 7, 2, 12, 1, 10, 14, 9},
                    {10, 6, 9, 0, 12, 11, 7, 13, 15, 1, 3, 14, 5, 2, 8, 4},
                    {3, 15, 0, 6, 10, 1, 13, 8, 9, 4, 5, 11, 12, 7, 2, 14}
            },
            {
                    {2, 12, 4, 1, 7, 10, 11, 6, 8, 5, 3, 15, 13, 0, 14, 9},
                    {14, 11, 2, 12, 4, 7, 13, 1, 5, 0, 15, 10, 3, 9, 8, 6},
                    {4, 2, 1, 11, 10, 13, 7, 8, 15, 9, 12, 5, 6, 3, 0, 14},
                    {11, 8, 12, 7, 1, 14, 2, 13, 6, 15, 0, 9, 10, 4, 5, 3}
            },
            {
                    {12, 1, 10, 15, 9, 2, 6, 8, 0, 13, 3, 4, 14, 7, 5, 11},
                    {10, 15, 4, 2, 7, 12, 9, 5, 6, 1, 13, 14, 0, 11, 3, 8},
                    {9, 14, 15, 5, 2, 8, 12, 3, 7, 0, 4, 10, 1, 13, 11, 6},
                    {4, 3, 2, 12, 9, 5, 15, 10, 11, 14, 1, 7, 6, 0, 8, 13}
            },
            {
                    {4, 11, 2, 14, 15, 0, 8, 13, 3, 12, 9, 7, 5, 10, 6, 1},
                    {13, 0, 11, 7, 4, 9, 1, 10, 14, 3, 5, 12, 2, 15, 8, 6},
                    {1, 4, 11, 13, 12, 3, 7, 14, 10, 15, 6, 8, 0, 5, 9, 2},
                    {6, 11, 13, 8, 1, 4, 10, 7, 9, 5, 0, 15, 14, 2, 3, 12}
            },
            {
                    {13, 2, 8, 4, 6, 15, 11, 1, 10, 9, 3, 14, 5, 0, 12, 7},
                    {1, 15, 13, 8, 10, 3, 7, 4, 12, 5, 6, 11, 0, 14, 9, 2},
                    {7, 11, 4, 1, 9, 12, 14, 2, 0, 6, 10, 13, 15, 3, 5, 8},
                    {2, 1, 14, 7, 4, 10, 8, 13, 15, 12, 9, 0, 3, 5, 6, 11}
            }
    };


    private final static int[] PBOX = new int[]{
            16, 7, 20, 21, 29, 12, 28, 17,
            1, 15, 23, 26, 5, 18, 31, 10,
            2, 8, 24, 14, 32, 27, 3, 9,
            19, 13, 30, 6, 22, 11, 4, 25
    };

    private final static int[] PC_1_C = new int[]{
            57, 49, 41, 33, 25, 17, 9,
            1, 58, 50, 42, 34, 26, 18,
            10, 2, 59, 51, 43, 35, 27,
            19, 11, 3, 60, 52, 44, 36
    };

    private final static int[] PC_1_D = new int[]{
            63, 55, 47, 39, 31, 23, 15,
            7, 62, 54, 46, 38, 30, 22,
            14, 6, 61, 53, 45, 37, 29,
            21, 13, 5, 28, 20, 12, 4
    };

    private final static int[] SHIFT = new int[]{
            1, 1, 2, 2, 2, 2, 2, 2,
            1, 2, 2, 2, 2, 2, 2, 1
    };

    private final static int[] PC_2 = new int[]{
            14, 17, 11, 24, 1, 5, 3, 28,
            15, 6, 21, 10, 23, 19, 12, 4,
            26, 8, 16, 7, 27, 20, 13, 2,
            41, 52, 31, 37, 47, 55, 30, 40,
            51, 45, 33, 48, 44, 49, 39, 56,
            34, 53, 46, 42, 50, 36, 29, 32
    };

    /**
     * 输入为32bit的右侧分组消息，输出为48bit的经过扩展置换的消息
     *
     * @param right 32 Bit的右侧分组消息
     * @return 48 Bit的经过扩展置换的消息
     * @throws Exception Bit位出错
     */
    private BitStream extend(BitStream right) throws Exception {
        // 经过扩展置换的消息，保存到result中
        BitStream result = new BitStream(48);
        for (int i = 0; i < 48; i++) {
            result.set(i + 1, right.get(EXTEND[i]));
        }
        return result;
    }

    /**
     * 将扩展后的消息与当前轮密钥进行异或运算
     *
     * @param afterExtend 扩展后的消息
     * @param roundKey    当前轮次的轮密钥
     * @return 返回一个处理后的消息
     */
    private BitStream addRoundKey(BitStream afterExtend, BitStream roundKey) throws Exception {
        for (int i = 1; i <= 48; i++) {
            afterExtend.set(i, afterExtend.get(i) ^ roundKey.get(i));
        }
        return afterExtend;
    }

    private BitStream sbox(BitStream afterAddRoundKey) throws Exception {
        // 用result来保存结果
        BitStream result = new BitStream(32);
        // 48bit分成8组，每组6bit，每一组使用对应的SBOXs进行处理
        for (int i = 0; i < 48; i += 6) {
            // 确定要使用的具体的SBOXs是哪一个
            int boxIndex = i / 6;
            // 确定SBOX中的元素的坐标
            int x = afterAddRoundKey.get(i + 1) * 2 + afterAddRoundKey.get(i + 6);
            int y = afterAddRoundKey.get(i + 2) * 8 + afterAddRoundKey.get(i + 3) * 4 +
                    afterAddRoundKey.get(i + 4) * 2 + afterAddRoundKey.get(i + 5);
            int resultUnit = SBOXs[boxIndex][x][y];
            result.set(boxIndex * 4 + 1, resultUnit / 8);
            result.set(boxIndex * 4 + 2, (resultUnit % 8) / 4);
            result.set(boxIndex * 4 + 3, (resultUnit % 4) / 2);
            result.set(boxIndex * 4 + 4, resultUnit % 2);
        }

        return result;
    }

    private BitStream pbox(BitStream afterSBOX) throws Exception {
        BitStream result = new BitStream(32);
        for (int i = 0; i < 32; i++) {
            result.set(i + 1, afterSBOX.get(PBOX[i]));
        }
        return result;
    }

    /**
     * 函数的输入是64位明文消息，输出是经过初始置换IP的64位消息
     *
     * @param message 64位明文消息
     * @return 64位IP初始置换后的消息
     */
    private BitStream IP(BitStream message) {
        BitStream result = null;
        try {
            result = new BitStream(message.getBitLength());
            for (int i = 0; i < 64; i++) {
                result.set(i + 1, message.get(IP[i]));
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }

        return result;
    }

    /**
     * 函数的输入为32位的右侧分组消息和轮密钥，输出为32位经过F函数的消息
     *
     * @param right 32位的右侧分组消息
     * @param K     某个轮密钥
     * @return 32位经过F函数的消息
     */
    private BitStream F(BitStream right, BitStream K) throws Exception {
        // 扩展置换 extend
        BitStream afterExtend = extend(right);
        // 和轮密钥按位异或，addRoundKey
        BitStream afterAddRoundKey = addRoundKey(afterExtend, K);
        // 输出经过S盒，sbox
        BitStream afterSBOX = sbox(afterAddRoundKey);
        // 线性置换，pbox
        BitStream afterPBOX = pbox(afterSBOX);
        return afterPBOX;
    }

    /**
     * 输入是经过合并的64位分组消息，输出是64位密文
     *
     * @param afterF 经过合并的64位分组消息
     * @return 64位密文
     */
    private BitStream inverseIP(BitStream afterF) throws Exception {
        // 使用ciphertext保存经过逆初始置换的消息
        BitStream ciphertext = new BitStream(afterF.getBitLength());
        for (int i = 0; i < 64; i++) {
            ciphertext.set(i + 1, afterF.get(IPInverse[i]));
        }
        return ciphertext;
    }


    // 密钥生成函数
    public BitStream[] generateKey(byte[] key) throws Exception {
        // 输入的是64比特（8字节）的原始密钥，但是每个字节的第8个比特位是奇偶校验位，需要去掉，所以只剩下56比特（7字节）
//        BitStream originK = new BitStream(64, key);
        // K是将key转化为本代码专用的比特流
        BitStream K = new BitStream(64, key);
//        int KIndex = 1;
//        for (int i = 1; i <= 64; i++) {
//            if (i % 8 == 0) {
//                continue;
//            }
//            K.set(KIndex++, originK.get(i));
//        }
        // 选择第1个置换矩阵进行置换，分为C、D两部分
        // C和D都是28个比特组成的数组
        BitStream C = new BitStream(28);
        BitStream D = new BitStream(28);
        for (int i = 0; i < 28; i++) {
            C.set(i + 1, K.get(PC_1_C[i]));
            D.set(i + 1, K.get(PC_1_D[i]));
        }

        // 16组48bit的轮密钥roundKey
        BitStream[] roundKeys = new BitStream[16];

        for (int i = 0; i < 16; i++) {
            roundKeys[i] = new BitStream(48);
            // 进行循环左移操作
            C.reLeftShift(SHIFT[i]);
            D.reLeftShift(SHIFT[i]);
            // 选择PC-2进行置换
            for (int j = 0; j < 48; j++) {
                if (PC_2[j] <= 28) {
                    roundKeys[i].set(j + 1, C.get(PC_2[j]));
                } else {
                    roundKeys[i].set(j + 1, D.get(PC_2[j] - 28));
                }
            }
        }

        return roundKeys;
    }

    /**
     * @param plaintext 64位的明文
     * @param roundKeys 16轮的轮密钥
     * @return 返回密文
     * @throws Exception
     */
    public BitStream encrypt(BitStream plaintext, BitStream[] roundKeys) throws Exception {
        // 初始置换IP(message)
        BitStream afterIP = IP(plaintext);
        // 将明文分成左半部分L[0]和右半部分R[0]
        BitStream left = new BitStream(32);
        for (int i = 1; i <= 32; i++) {
            left.set(i, afterIP.get(i));
        }
        BitStream right = new BitStream(32);
        for (int i = 1; i <= 32; i++) {
            right.set(i, afterIP.get(32 + i));
        }
        // 轮函数加密迭代，一共进行16轮
        BitStream temp = new BitStream(32);
        for (int i = 0; i < 16; i++) {
            // L[i] = R[i-1];
            for (int index = 1; index <= 32; index++) {
                temp.set(index, right.get(index));
            }
            // R[i]=xor(L[i-1],f(R[i-1],roundKey[i]));
            BitStream f = F(right, roundKeys[i]);
            for (int index = 1; index <= 32; index++) {
                right.set(index, left.get(index) ^ f.get(index));
            }
            // left = temp;
            for (int index = 1; index <= 32; index++) {
                left.set(index, temp.get(index));
            }
        }
        // 左、右部分合并，注意是RL，不是LR
        BitStream afterF = new BitStream(64);
        for (int index = 1; index <= 32; index++) {
            afterF.set(index + 32, left.get(index));
            afterF.set(index, right.get(index));
        }
        // 逆初始置换inverseIP(afterF)
        BitStream result = inverseIP(afterF);

        return result;
    }


    public BitStream decrypt(BitStream ciphertext, BitStream[] roundKeys) throws Exception {
        // 初始置换IP(message)
        BitStream afterIP = IP(ciphertext);
        // 将明文分成左半部分L[0]和右半部分R[0]
        BitStream left = new BitStream(32);
        for (int i = 1; i <= 32; i++) {
            left.set(i, afterIP.get(i));
        }
        BitStream right = new BitStream(32);
        for (int i = 1; i <= 32; i++) {
            right.set(i, afterIP.get(32 + i));
        }
        // 轮函数加密迭代，一共进行16轮
        BitStream temp = new BitStream(32);
        for (int i = 0; i < 16; i++) {
            // L[i] = R[i-1];
            for (int index = 1; index <= 32; index++) {
                temp.set(index, right.get(index));
            }
            // R[i]=xor(L[i-1],f(R[i-1],roundKey[i]));
            BitStream f = F(right, roundKeys[15 - i]);
            for (int index = 1; index <= 32; index++) {
                right.set(index, left.get(index) ^ f.get(index));
            }
            // left = temp;
            for (int index = 1; index <= 32; index++) {
                left.set(index, temp.get(index));
            }
        }
        // 左、右部分合并，注意是RL，不是LR
        BitStream afterF = new BitStream(64);
        for (int index = 1; index <= 32; index++) {
            afterF.set(index + 32, left.get(index));
            afterF.set(index, right.get(index));
        }
        // 逆初始置换inverseIP(afterF)
        BitStream plaintext = inverseIP(afterF);

        return plaintext;
    }

}
